package com.ideaaedi.component.decompile.cfr;

import com.ideaaedi.component.decompile.Decompiler;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.benf.cfr.reader.api.CfrDriver;
import org.benf.cfr.reader.api.OutputSinkFactory;
import org.benf.cfr.reader.util.AnalysisType;

import java.io.CharArrayReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.TreeMap;

/**
 * cfr反编译器<br><br>
 * 注：CFR反编译器，在泛型、invokedynamic等的处理上不如fernflower强大。
 *
 * @author JustryDeng
 * @since 2021/9/23 22:59:27
 */
public class CfrDecompiler implements Decompiler {
    
    /**
     * 反编译
     *
     * @param classByteCollection
     *            待反编译的字节码集合
     * @param methodName
     *            只反编译指定的方法
     * @return  反编译出来的<b>源码的classFilePath</b>以及<b>源码</b>.
     *          注意：<ul>
     *                  <li>key - 源码对应class的classFilePath, 形如com/szlzcl/jdk8feature/OptionalTests$MyConsumer.class<br/>详见{@link CfrDecompiler#obtainClassFilePath(byte[])}</li>
     *                  <li>value - 源码<br>注意：反编译失败或找不存要反编译的目标方法时，值为空白字符串</li>
     *               </ul>
     */
    @Override
    public LinkedHashMap<String, String> decompile(Collection<byte[]> classByteCollection, String methodName) {
        if (classByteCollection == null || classByteCollection.size() == 0) {
            throw new IllegalArgumentException("classByteCollection cannot be empty.");
        }
        // 可选项
        HashMap<String, String> options = new HashMap<>(4);
        // 避免中文乱码
        options.put("hideutf", "false");
        // 反编译出来的java前面，没必要打印所使用CFR的版本号
        options.put("showversion", "false");
        // 批量反编译时，支持反编译内部类
        options.put("skipbatchinnerclasses", "false");
        // 反编译出来的java头部，不需要打印CFR相关的注释
        options.put("comments", "false");
        // 设置方便易目标是.class
        options.put("analyseas", AnalysisType.CLASS.name());
        // 只反编译指定的类
        if (StringUtils.isNotBlank(methodName)) {
            options.put("methodname", methodName);
        }
        /*
         * 需要反编译的字节码容器
         * <br>
         * 根据CfrDriver的实现CfrDriverImpl, CfrDriverImpl#analyse(classFilePathList)方法内部会对
         * 传入的classFilePathList集合进行自然排序，然后再挨个进行反编译。
         * 这里在外层就主动自然排序一下，使classFilePathList中元素顺序与反编
         * 译出来的sourceCodeList中的元素能一一对应
         */
        NavigableMap<String, byte[]> classProvider = new TreeMap<>();
        for (byte[] classBytes : classByteCollection) {
            String classFilePath = obtainClassFilePath(classBytes);
            classProvider.put(classFilePath, classBytes);
        }
        
        // 反编译出来的源码容器
        List<String> sourceCodeList = new ArrayList<>(2);
        // 构造反编译驱动
        CfrDriver driver = new CfrDriver.Builder()
                .withOptions(options)
                .withClassFileSource(new ClassBytesSource(classProvider))
                .withOutputSink(new SourceCodeOutputSinkFactory(sourceCodeList))
                .build();
        // 执行反编译
        List<String> classFilePathList = new ArrayList<>(classProvider.keySet());
        // 提示：能否反编译内部类成功，classFilePathArray中item的名字很关键
        driver.analyse(classFilePathList);
    
        LinkedHashMap<String, String> result = new LinkedHashMap<>();
        for (int i = 0; i < classFilePathList.size(); i++) {
            result.put(classFilePathList.get(i), sourceCodeList.get(i));
        }
        return result;
    }
    
    /**
     * 获取classpath下，.class文件对应的路径名
     *
     * @param classBytes
     *            类字节码
     * @return classpath下，.class文件对应的路径名. <br> 形如: com/szlzcl/jdk8feature/OptionalTests$MyConsumer.class
     */
    private static String obtainClassFilePath(byte[] classBytes) {
        if (classBytes == null) {
            throw new IllegalArgumentException("classBytes cannot be empty.");
        }
        StringBuilder sb = new StringBuilder();
        Map<String, byte[]> classProvider = new HashMap<>(2);
        classProvider.put("Tmp_" + System.currentTimeMillis() + ".class", classBytes);
        CfrDriver driver = new CfrDriver.Builder()
                .withClassFileSource(new ClassBytesSource(classProvider))
                .withOutputSink(new OutputSinkFactory() {
    
                    @Override
                    public List<SinkClass> getSupportedSinks(SinkType sinkType, Collection<SinkClass> available) {
                        return Arrays.asList(SinkClass.DECOMPILED, SinkClass.DECOMPILED_MULTIVER);
                    }
    
                    @Override
                    public <T> Sink<T> getSink(SinkType sinkType, SinkClass sinkClass) {
                        return sinkable -> {
                            // 反编译进度信息不需要记录
                            if (sinkType == SinkType.PROGRESS && sinkable != null) {
                                sb.append("\n// ").append(sinkable);
                            }
                        };
                    }
                })
                .build();
        // 执行反编译
        driver.analyse(new ArrayList<>(classProvider.keySet()));
    
        List<String> lines;
        try {
            lines = IOUtils.readLines(new CharArrayReader(sb.toString().toCharArray()));
        } catch (IOException e) {
            throw  new DecompileException("Execute code IOUtils.readLines(...) exception.", e);
        }
        String flag = "Analysing type";
        for (String line : lines) {
            // line形如：// Analysing type com.szlzcl.jdk8feature.OptionalTests$MyConsumer
            int idx = line.indexOf(flag);
            if (idx >= 0) {
                // 得到: com.szlzcl.jdk8feature.OptionalTests$MyConsumer
                String likeClassLongName = line.substring(idx + flag.length()).trim();
                // 得到: com/szlzcl/jdk8feature/OptionalTests$MyConsumer.class
                return likeClassLongName.replace(".", "/") + ".class";
            }
        }
        throw new DecompileException("Cannot obtain classFilePath from lines -> " + lines);
    }
}
